/*
 * tcpsession.c
 *
 *  Created on: 2012-8-2
 *      Author: wx
 */
#include <stdio.h>
#include <stdlib.h>
#include <strings.h>
#include <string.h>
#include <arpa/inet.h>
#include <byteswap.h>

#include "tcpsession.h"
#include "util.h"
#include "log.h"

#define MAX_SESSION_COUNT 1048576

static TCPSession_t **m_TCPSessions;
static struct pcap_file_header m_FileHeader;
static int m_SessionCount = 0;


static void ntohheader(FramHeader_t * frame_header, IPHeader_t * ip_header, TCPHeader_t * tcp_header);
static TCPSession_t * GetSession(u_int32 srcip, u_int16 srcport, u_int32 dstip, u_int16 dstport);
//static int CheckSession(int index);

//初始化
int tcpsess_init()
{
	//分配会话数据存储资源
	m_TCPSessions = malloc(sizeof(TCPSession_t *) * MAX_SESSION_COUNT);
	bzero(m_TCPSessions, sizeof(TCPSession_t *) * MAX_SESSION_COUNT);
	m_SessionCount = 0;
	return 0;
}

int tcpsess_parse_pcap(const char * pcap_filename)
{
	static FILE *fp;
	static struct pcap_pkthdr pkt_header;
	static FramHeader_t frame_header;
	static IPHeader_t ip_header;
	static TCPHeader_t tcp_header;
	static TCPSession_t *session;
	static int ip_header_len, tcp_header_len;
	static int pos;
	static int pkt_count=0;
	static int tcp_data_len;
	static char tcp_data[2048];


	if ((fp = fopen(pcap_filename, "r")) == NULL)
	{
		SendLog("Failed to open pcap file.");
		return TSERR_FILEOPEN;
	}

	pos = 0;
	if (fread(&m_FileHeader, sizeof(m_FileHeader), 1, fp) != 1 )
	{
		SendLog("Failed to read pcap file.");
		fclose(fp);
		return TSERR_FILEREAD;
	}
	pos += sizeof(m_FileHeader);

	//判断字节顺序
	if (m_FileHeader.magic == 0xd4c3b2a1)
	{
		//单独转换文件头的字节顺序
		m_FileHeader.version_major = bswap_16(m_FileHeader.version_major);
		m_FileHeader.version_minor = bswap_16(m_FileHeader.version_minor);
		m_FileHeader.thiszone = bswap_32(m_FileHeader.thiszone);
		m_FileHeader.sigfigs = bswap_32(m_FileHeader.sigfigs);
		m_FileHeader.snaplen = bswap_32(m_FileHeader.snaplen);
		m_FileHeader.linktype = bswap_32(m_FileHeader.linktype);
	} else if (m_FileHeader.magic != 0xa1b2c3d4)
	{
		SendLog("Unsupported file format");
		fclose(fp);
		return TSERR_UNSUPPORTED;
	}

	//如果是以太网，继续。
	if (m_FileHeader.linktype != 1)
	{
		SendLog("Unsupported link type.");
		fclose(fp);
		return TSERR_UNSUPPORTED;
	}

	while (fread(&pkt_header, sizeof(pkt_header), 1, fp) == 1)
	{
		pkt_count++;
//SendLog("Packet Number: %d\t", pkt_count);
		pos += sizeof(pkt_header);

//SendLog("%s\t", GetTimeStr(pkt_header.ts.tv_sec, pkt_header.ts.tv_usec));
		if (m_FileHeader.magic == 0xd4c3b2a1)
		{
			pkt_header.caplen = bswap_32(pkt_header.caplen);
			pkt_header.len = bswap_32(pkt_header.len);
			pkt_header.ts.tv_sec = bswap_32(pkt_header.ts.tv_sec);
			pkt_header.ts.tv_usec = bswap_32(pkt_header.ts.tv_usec);
		}

		if (fread(&frame_header, sizeof(frame_header), 1, fp) != 1 )
		{
			SendLog("Failed to read pcap file.");
			fclose(fp);
			return TSERR_FILEREAD;
		}
		pos += sizeof(frame_header);

		if (frame_header.FrameType != htons(0x0800))
		{
			//丢掉数据包
			if (fseek(fp, pkt_header.caplen - sizeof(frame_header), SEEK_CUR) == -1 )
			{
				SendLog("Failed to read pcap file.");
				fclose(fp);
				return TSERR_FILEREAD;
			}
			pos += pkt_header.caplen - sizeof(frame_header);
			continue;
		}

		if (fread(&ip_header, sizeof(ip_header), 1, fp) !=1 )
		{
			SendLog("Failed to read pcap file.");
			fclose(fp);
			return TSERR_FILEREAD;
		}
		pos += sizeof(ip_header);

//SendLog("Src IP=%s", GetIPStr(ip_header.SrcIP));
//SendLog(", Dst IP=%s\t\t", GetIPStr(ip_header.DstIP));

		if (ip_header.Protocol != 6)
		{
			//丢掉数据包
			if (fseek(fp, pkt_header.caplen - sizeof(frame_header) - sizeof(ip_header), SEEK_CUR) == -1 )
			{
				SendLog("Failed to read pcap file.");
				fclose(fp);
				return TSERR_FILEREAD;
			}
			pos += pkt_header.caplen - sizeof(frame_header) - sizeof(ip_header);
			continue;
		}

		//根据包头长度判断ip头是否有可选项，文件指针跳过可选项。
		ip_header_len = (ip_header.Ver_HLen & 0x0F) * 4;
		if ((ip_header_len - sizeof(ip_header)) > 0)
		{
			if (fseek(fp, ip_header_len - sizeof(ip_header), SEEK_CUR) == -1)
			{
				SendLog("Failed to read pcap file.");
				fclose(fp);
				return TSERR_FILEREAD;
			}
			pos += ip_header_len - sizeof(ip_header);
		}

		if (fread(&tcp_header, sizeof(tcp_header), 1, fp) != 1 )
		{
			SendLog("Failed to read pcap file.");
			fclose(fp);
			return TSERR_FILEREAD;
		}
		pos += sizeof(tcp_header);

//SendLog("Src Port=%d, Dst Port=%d\t\t", ntohs(tcp_header.SrcPort), ntohs(tcp_header.DstPort));

		//跳过tcp头可选项
		tcp_header_len = ((tcp_header.HeaderLen & 0xF0)>>4)<<2;
		if ((tcp_header_len - sizeof(tcp_header))>0)
		{
			if (fseek(fp, tcp_header_len - sizeof(tcp_header), SEEK_CUR) == -1)
			{
				SendLog("Failed to read pcap file.");
				fclose(fp);
				return TSERR_FILEREAD;
			}
			pos += tcp_header_len - sizeof(tcp_header);
		}

		//各个头网络字节顺序转换
		ntohheader(&frame_header, &ip_header, &tcp_header);

		//tcp载荷数据长度
		tcp_data_len = ip_header.TotalLen - ip_header_len - tcp_header_len;
		if (tcp_data_len>0)
		{
			//读取数据
			if (fread(tcp_data, tcp_data_len, 1, fp) != 1)
			{
				SendLog("Failed to read pcap file.");
				fclose(fp);
				return TSERR_FILEREAD;
			}
			pos += tcp_data_len;
		}

		//判断是否有填充数据，读取tcp数据
		if (ip_header.TotalLen<46)
		{
			if (fseek(fp, 46 - ip_header.TotalLen, SEEK_CUR) == -1)
			{
				SendLog("Failed to read pcap file.");
				fclose(fp);
				return TSERR_FILEREAD;
			}
			pos += 46 - ip_header.TotalLen;
		}

		//处理数据包
/*		if (tcp_header.Flags & TCPFLAG_SYN) SendLog("SYN ");
		if (tcp_header.Flags & TCPFLAG_ACK) SendLog("ACK ");
		if (tcp_header.Flags & TCPFLAG_PSH) SendLog("PSH ");
		if (tcp_header.Flags & TCPFLAG_FIN) SendLog("FIN ");
		if (tcp_header.Flags & TCPFLAG_RST) SendLog("RST ");
		if (tcp_header.Flags & TCPFLAG_UTG) SendLog("UTG ");
*/
		session = GetSession(ip_header.SrcIP, tcp_header.SrcPort, ip_header.DstIP, tcp_header.DstPort);
		if (session)
		{
			//找到session
			while (session->next != NULL) session=session->next;
			session->next = malloc(sizeof(TCPSession_t));
			session = session->next;
		} else
		{
			//创建新的session
			*(m_TCPSessions + m_SessionCount) = malloc(sizeof(TCPSession_t));
			session = *(m_TCPSessions + m_SessionCount);
			m_SessionCount++;
		}
		bzero(session, sizeof(TCPSession_t));
		session->PkgNo = pkt_count;
		session->ts = pkt_header.ts;
		memcpy(session->DstMAC, frame_header.DstMAC, 6);
		memcpy(session->SrcMAC, frame_header.SrcMAC, 6);
		session->TotalLen = ip_header.TotalLen;
		session->ID = ip_header.ID;
		session->SrcIP = ip_header.SrcIP;
		session->DstIP = ip_header.DstIP;
		session->SrcPort = tcp_header.SrcPort;
		session->DstPort = tcp_header.DstPort;
		session->SeqNo = tcp_header.SeqNO;
		session->AckNo = tcp_header.AckNO;
		session->Flags = tcp_header.Flags;
		if (tcp_data_len>0)
		{
			session->DataLen = tcp_data_len;
			session->Payload = malloc(tcp_data_len);
			memcpy(session->Payload, tcp_data, tcp_data_len);
		}



/*
		if (tcp_header.Flags == TCPFLAG_SYN)
		{
			//第一次握手
			//创建会话
SendLog("SYN only ");

		} else if (tcp_header.Flags == (TCPFLAG_SYN||TCPFLAG_ACK))
		{
			//第二次握手
		} else if (tcp_header.Flags == TCPFLAG_ACK)
		{
			//第三次握手及以后应答
		}
*/

//SendLog("\n");
	}

//	for (i=0;i<m_SessionCount;i++) CheckSession(i);

//SendLog("Session Count: %d\n", m_SessionCount);


	fclose(fp);

	return m_SessionCount;
}

//释放解析后的数据所占用的内存
void tcpsess_close()
{
	static int i;
	static TCPSession_t *session, *next;
	for (i=0;i<MAX_SESSION_COUNT;i++)
	{
		session = *(m_TCPSessions+i);
		if (session)
		{
			while(session)
			{
				next = session->next;
				if (session->Payload) free(session->Payload);
				free(session);
				session=next;
			}
		}
	}

	free(m_TCPSessions);
	m_SessionCount = 0;
}

struct pcap_file_header * tcpsess_get_cap_header()
{
	return &m_FileHeader;
}

//将各个头从网络字节顺序转换为主机字节顺序
//ip地址不转换
static void ntohheader(FramHeader_t *frame_header, IPHeader_t *ip_header, TCPHeader_t *tcp_header)
{
	if (frame_header != NULL)
	{
		frame_header->FrameType = ntohs(frame_header->FrameType);
	}

	if (ip_header != NULL)
	{
		ip_header->TotalLen = ntohs(ip_header->TotalLen);
		ip_header->ID = ntohs(ip_header->ID);
		ip_header->Flag_Segment = ntohs(ip_header->Flag_Segment);
		ip_header->Checksum = ntohs(ip_header->Checksum);
		//ip_header->SrcIP = ntohl(ip_header->SrcIP);
		//ip_header->DstIP = ntohl(ip_header->DstIP);
	}

	if (tcp_header != NULL)
	{
		tcp_header->SrcPort = ntohs(tcp_header->SrcPort);
		tcp_header->DstPort = ntohs(tcp_header->DstPort);
		tcp_header->SeqNO = ntohl(tcp_header->SeqNO);
		tcp_header->AckNO = ntohl(tcp_header->AckNO);
		tcp_header->Window = ntohs(tcp_header->Window);
		tcp_header->Checksum = ntohs(tcp_header->Checksum);
		tcp_header->UrgentPointer = ntohs(tcp_header->UrgentPointer);
	}
}

//根据地址和端口查找会话，如果没找到，返回NULL
static TCPSession_t * GetSession(u_int32 srcip, u_int16 srcport, u_int32 dstip, u_int16 dstport)
{
	static int i;
	static TCPSession_t *session;

	for (i=0;i<m_SessionCount;i++)
	{
		session = *(m_TCPSessions+i);
		if (((session->SrcIP==srcip)&&(session->DstIP==dstip)&&(session->SrcPort==srcport)&&(session->DstPort==dstport))
		  ||((session->SrcIP==dstip)&&(session->DstIP==srcip)&&(session->SrcPort==dstport)&&(session->DstPort==srcport)))
		{
			return session;
		}
	}
	return NULL;
}

//检查指定session是否是完整的会话过程
//static int CheckSession(int index)
//{
//	//检查三次握手，seq、ack顺序，检查fin或rst，删除重发包
//	return 0;
//}

//获取会话数组
TCPSession_t ** tcpsess_get_sessions()
{
	return m_TCPSessions;
}

//获取会话数量
int tcpsess_get_sess_count()
{
	return m_SessionCount;
}

//删除指定session
int tcpsess_del_session(int index)
{
	static int i;
	static TCPSession_t *session, *next;

	//删除会话后，要把数组最后一个移动过来补位。以后可以考虑采用链表结构。

	if ((index>m_SessionCount-1)||(index<0)) return TSERR_OUTOFINDEX;
	session = *(m_TCPSessions + index);
	if (session)
	{
		while(session)
		{
			next = session->next;
			if (session->Payload) free(session->Payload);
			free(session);
			session=next;
		}
	}

	if (index != m_SessionCount-1)
	{
		for (i=index;i<m_SessionCount-1;i++)
		{
			*(m_TCPSessions + i) = *(m_TCPSessions + i + 1);
		}
	}
	*(m_TCPSessions + m_SessionCount-1) = NULL;

	m_SessionCount--;


	return 0;
}


//按照时序获取所有会话数据。
//成功返回数据长度，失败返回错误代码。
int tcpsess_get_sess_data(TCPSession_t *session, char * buf, int bufsize)
{
	//整理会话数据包，过滤掉错误、重发数据包，得到正确的会话数据
	static TCPSession_t * s;
	static int datalen;
	static u_int32 sseq, dseq;
	static u_int32 ipa, ipb;
	static u_int16 porta, portb;

	if (!session) return TSERR_UNKNOWN;
	datalen = 0;
	s = session;
	ipa = s->SrcIP;
	porta = s->SrcPort;
	ipb = s->DstIP;
	portb = s->DstPort;
	sseq = 0;
	dseq = 0;
	bzero(buf, bufsize);
	while (s)
	{
		//根据seqno检查数据包是否有错
		if ((s->SrcIP == ipa)&&(s->SrcPort == porta))
		{
			if (sseq == 0)
			{
				sseq = s->SeqNo;
				if ((s->Flags & TCPFLAG_SYN) || (s->Flags & TCPFLAG_FIN))
				{
					sseq++;
				} else
				{
					sseq += s->DataLen;
				}
			} else
			{
				if (sseq != s->SeqNo)
				{
					//丢掉当前包
					s = s->next;
					continue;
				}
				sseq += s->DataLen;
			}
		} else
		{
			if (dseq == 0)
			{
				dseq = s->SeqNo;
				if ((s->Flags & TCPFLAG_SYN) || (s->Flags & TCPFLAG_FIN))
				{
					dseq++;
				} else
				{
					dseq += s->DataLen;
				}
			} else
			{
				if (dseq != s->SeqNo)
				{
					//丢掉当前包
					s = s->next;
					continue;
				}
				dseq += s->DataLen;
			}
		}
		if (s->DataLen>0)
		{
			datalen += s->DataLen;
			if (datalen>bufsize) return TSERR_OVERFLOW;
			memcpy(buf + datalen - s->DataLen, s->Payload, s->DataLen);
		}
		s = s->next;
	}

	return datalen;

}

//返回所有请求数据
//成功返回数据长度，失败返回错误代码。
int tcpsess_get_req_data(TCPSession_t *session, char * buf, int bufsize)
{
	//整理会话数据包，过滤掉错误、重发数据包，得到正确的会话数据
	static TCPSession_t * s;
	static int datalen;
	static u_int32 sseq;
	static u_int32 ipa;
	static u_int16 porta;

	if (!session) return TSERR_UNKNOWN;
	datalen = 0;
	s = session;
	ipa = s->SrcIP;
	porta = s->SrcPort;
	sseq = 0;
	bzero(buf, bufsize);
	while (s)
	{
		//根据seqno检查数据包是否有错
		if ((s->SrcIP == ipa)&&(s->SrcPort == porta))
		{
			if (sseq == 0)
			{
				sseq = s->SeqNo;
				if ((s->Flags & TCPFLAG_SYN) || (s->Flags & TCPFLAG_FIN))
				{
					sseq++;
				} else
				{
					sseq += s->DataLen;
				}
			} else
			{
				if (sseq != s->SeqNo)
				{
					//丢掉当前包
					s = s->next;
					continue;
				}
				sseq += s->DataLen;
			}
			if (s->DataLen>0)
			{
				datalen += s->DataLen;
				if (datalen>bufsize) return TSERR_OVERFLOW;
				memcpy(buf + datalen - s->DataLen, s->Payload, s->DataLen);
			}
		}
		s = s->next;
	}

	return datalen;
}

//返回所有应答数据
//成功返回数据长度，失败返回错误代码。
int tcpsess_get_rsp_data(TCPSession_t *session, char * buf, int bufsize)
{
	//整理会话数据包，过滤掉错误、重发数据包，得到正确的会话数据
	static TCPSession_t * s;
	static int datalen;
	static u_int32 sseq;
	static u_int32 ipa;
	static u_int16 porta;

	if (!session) return TSERR_UNKNOWN;
	datalen = 0;
	s = session;
	ipa = s->DstIP;
	porta = s->DstPort;
	sseq = 0;
	bzero(buf, bufsize);
	while (s)
	{
		//根据seqno检查数据包是否有错
		if ((s->SrcIP == ipa)&&(s->SrcPort == porta))
		{
			if (sseq == 0)
			{
				sseq = s->SeqNo;
				if ((s->Flags & TCPFLAG_SYN) || (s->Flags & TCPFLAG_FIN))
				{
					sseq++;
				} else
				{
					sseq += s->DataLen;
				}
			} else
			{
				if (sseq != s->SeqNo)
				{
					//丢掉当前包
					s = s->next;
					continue;
				}
				sseq += s->DataLen;
			}
			if (s->DataLen>0)
			{
				datalen += s->DataLen;
				if (datalen>bufsize) return TSERR_OVERFLOW;
				memcpy(buf + datalen - s->DataLen, s->Payload, s->DataLen);
			}
		}
		s = s->next;
	}

	return datalen;
}
